# coding=utf-8

import numpy as np


class ProductionEnv:
    def __init__(self):
        self.state = np.zeros(3)  # 初始状态：三个工厂的生产量
        self.action_space = ['increase_x', 'decrease_x', 'increase_y', 'decrease_y']
        self.base_cost = np.array([20, 25, 30])  # 基础成本
        self.base_quality = np.array([0.9, 0.85, 0.8])  # 初始合格率
        self.cost_multiplier = 1.2  # 合格率提高导致成本增加的非线性倍数

    def reset(self):
        self.state = np.zeros(3)
        return self.state

    def step(self, action):
        # 当前合格率和成本
        current_quality = self.base_quality.copy()
        current_cost = self.base_cost.copy()

        if action == 'increase_x':
            self.state += 10  # 增加原材料投入
        elif action == 'decrease_x':
            self.state = np.maximum(0, self.state - 10)  # 减少原材料投入

        elif action == 'increase_y':
            for i in range(len(current_quality)):
                current_quality[i] = min(1.0, current_quality[i] + 0.1)  # 增加合格率
                current_cost[i] += self.cost_multiplier * (current_quality[i] - self.base_quality[i]) ** 2  # 非线性增加成本
        elif action == 'decrease_y':
            for i in range(len(current_quality)):
                current_quality[i] = max(0.0, current_quality[i] - 0.1)  # 减少合格率
                current_cost[i] -= self.cost_multiplier * (self.base_quality[i] - current_quality[i]) ** 2  # 成本适度减少
                current_cost[i] = max(0, current_cost[i])  # 确保成本不为负

        # 计算奖励（利润）
        reward = self.calculate_profit(current_quality, current_cost)

        return self.state, reward

    def calculate_profit(self, quality, cost):
        # 计算总利润
        prices = np.array([100, 120, 130])  # 产品价格
        profits = (prices - cost) * quality * self.state  # 利润 = (价格 - 成本) * 合格率 * 产量
        return np.sum(profits)


# 强化学习代理训练示例
env = ProductionEnv()
state = env.reset()

for episode in range(1000):  # 多次迭代训练
    action = np.random.choice(env.action_space)  # 随机选择动作
    next_state, reward = env.step(action)  # 获取下一状态和奖励
    print(action, reward, next_state)
